#!/bin/bash
# pg_dump 恢复脚本

# 脚本目录路径
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
COMMON_DIR="${SCRIPT_DIR}/../common"

# 加载公共工具函数
source "${COMMON_DIR}/utils.sh"

# 显示使用说明
show_usage() {
    echo "用法: $0 [选项] BACKUP_FILE"
    echo "选项:"
    echo "  --db=DB_NAME        目标数据库名称，默认为 ${POSTGRES_DB}"
    echo "  --create            如果数据库不存在，则创建它"
    echo "  --drop              恢复前删除数据库（如果存在）"
    echo "  --clean             添加 --clean 选项以覆盖现有对象"
    echo "  --verbose           详细输出模式"
    echo "  --help              显示此帮助信息"
    echo
    echo "示例:"
    echo "  $0 /path/to/backup.dump                # 恢复指定备份到默认数据库"
    echo "  $0 --db=newdb --create backup.dump     # 恢复到新数据库（如果不存在则创建）"
    echo "  $0 --db=mydb --drop --clean backup.dump # 删除现有数据库并重新创建"
}

# 初始化环境
LOG_FILE_PREFIX="pg_dump_restore"
init
start_timer

# 解析命令行参数
DB_NAME="${POSTGRES_DB}"
CREATE_DB=false
DROP_DB=false
CLEAN=false
VERBOSE=false
BACKUP_FILE=""

# 处理参数
while [ $# -gt 0 ]; do
    case "$1" in
        --db=*)
            DB_NAME="${1#*=}"
            shift
            ;;
        --create)
            CREATE_DB=true
            shift
            ;;
        --drop)
            DROP_DB=true
            shift
            ;;
        --clean)
            CLEAN=true
            shift
            ;;
        --verbose)
            VERBOSE=true
            shift
            ;;
        --help)
            show_usage
            exit 0
            ;;
        -*)
            log_error "未知选项: $1"
            show_usage
            exit 1
            ;;
        *)
            if [ -z "${BACKUP_FILE}" ]; then
                BACKUP_FILE="$1"
            else
                log_error "多余的参数: $1"
                show_usage
                exit 1
            fi
            shift
            ;;
    esac
done

# 验证必需的参数
if [ -z "${BACKUP_FILE}" ]; then
    log_error "必须指定备份文件"
    show_usage
    exit 1
fi

# 检查备份文件是否存在
if [ ! -e "${BACKUP_FILE}" ]; then
    log_error "备份文件不存在: ${BACKUP_FILE}"
    exit 1
fi

log_info "开始恢复数据库 ${DB_NAME} 从备份: ${BACKUP_FILE}"

# 检查 PostgreSQL 连接
if ! check_postgres_connection; then
    log_error "PostgreSQL 连接检查失败，无法继续恢复"
    exit 1
fi

# 设置 PostgreSQL 密码环境变量
export PGPASSWORD="${POSTGRES_PASSWORD}"

# 检查数据库是否存在
DB_EXISTS=false
if psql -h "${POSTGRES_HOST}" -p "${POSTGRES_PORT}" -U "${POSTGRES_USER}" -lqt | cut -d \| -f 1 | grep -qw "${DB_NAME}"; then
    DB_EXISTS=true
    log_info "数据库 ${DB_NAME} 已存在"

    # 如果指定了删除，则删除数据库
    if [ "${DROP_DB}" = "true" ]; then
        log_warning "将删除现有数据库 ${DB_NAME}"
        if ! psql -h "${POSTGRES_HOST}" -p "${POSTGRES_PORT}" -U "${POSTGRES_USER}" -c "DROP DATABASE \"${DB_NAME}\""; then
            log_error "无法删除数据库 ${DB_NAME}"
            unset PGPASSWORD
            exit 1
        fi
        DB_EXISTS=false
        log_success "数据库 ${DB_NAME} 已删除"
    fi
else
    log_info "数据库 ${DB_NAME} 不存在"
fi

# 如果数据库不存在且需要创建
if [ "${DB_EXISTS}" = "false" ] && [ "${CREATE_DB}" = "true" -o "${DROP_DB}" = "true" ]; then
    log_info "创建数据库 ${DB_NAME}"
    if ! psql -h "${POSTGRES_HOST}" -p "${POSTGRES_PORT}" -U "${POSTGRES_USER}" -c "CREATE DATABASE \"${DB_NAME}\""; then
        log_error "无法创建数据库 ${DB_NAME}"
        unset PGPASSWORD
        exit 1
    fi
    log_success "数据库 ${DB_NAME} 已创建"
elif [ "${DB_EXISTS}" = "false" ] && [ "${CREATE_DB}" = "false" ]; then
    log_error "数据库 ${DB_NAME} 不存在，使用 --create 选项自动创建它"
    unset PGPASSWORD
    exit 1
fi

# 检查备份文件类型并选择合适的恢复命令
BACKUP_FORMAT=""
RESTORE_CMD=""

if [ -d "${BACKUP_FILE}" ]; then
    # 目录格式备份
    BACKUP_FORMAT="directory"
    RESTORE_CMD="pg_restore -h ${POSTGRES_HOST} -p ${POSTGRES_PORT} -U ${POSTGRES_USER} -d \"${DB_NAME}\""
elif [[ "${BACKUP_FILE}" == *.dump ]]; then
    # Custom 格式备份
    BACKUP_FORMAT="custom"
    RESTORE_CMD="pg_restore -h ${POSTGRES_HOST} -p ${POSTGRES_PORT} -U ${POSTGRES_USER} -d \"${DB_NAME}\""
elif [[ "${BACKUP_FILE}" == *.tar ]]; then
    # Tar 格式备份
    BACKUP_FORMAT="tar"
    RESTORE_CMD="pg_restore -h ${POSTGRES_HOST} -p ${POSTGRES_PORT} -U ${POSTGRES_USER} -d \"${DB_NAME}\""
elif [[ "${BACKUP_FILE}" == *.sql ]] || [[ "${BACKUP_FILE}" == *.sql.gz ]]; then
    # Plain 格式备份
    BACKUP_FORMAT="plain"

    if [[ "${BACKUP_FILE}" == *.gz ]]; then
        # 如果是压缩的，需要解压
        RESTORE_CMD="gunzip -c \"${BACKUP_FILE}\" | psql -h ${POSTGRES_HOST} -p ${POSTGRES_PORT} -U ${POSTGRES_USER} -d \"${DB_NAME}\""
    else
        RESTORE_CMD="psql -h ${POSTGRES_HOST} -p ${POSTGRES_PORT} -U ${POSTGRES_USER} -d \"${DB_NAME}\" -f \"${BACKUP_FILE}\""
    fi
else
    # 尝试自动检测格式
    log_warning "无法从文件扩展名确定备份格式，尝试使用 pg_restore"
    BACKUP_FORMAT="unknown"
    RESTORE_CMD="pg_restore -h ${POSTGRES_HOST} -p ${POSTGRES_PORT} -U ${POSTGRES_USER} -d \"${DB_NAME}\""
fi

# 添加选项
if [ "${CLEAN}" = "true" ] && [ "${BACKUP_FORMAT}" != "plain" ]; then
    RESTORE_CMD="${RESTORE_CMD} --clean"
fi

if [ "${VERBOSE}" = "true" ]; then
    RESTORE_CMD="${RESTORE_CMD} -v"
fi

# 如果不是 plain 格式，添加备份文件路径
if [ "${BACKUP_FORMAT}" != "plain" ] || [[ "${BACKUP_FILE}" != *.gz ]]; then
    RESTORE_CMD="${RESTORE_CMD} \"${BACKUP_FILE}\""
fi

# 显示将要执行的命令
log_info "恢复命令: ${RESTORE_CMD}"

# 执行恢复
log_info "开始恢复数据库 ${DB_NAME}..."
if eval "${RESTORE_CMD}"; then
    log_success "数据库恢复成功完成"
else
    RESTORE_EXIT_CODE=$?
    log_error "数据库恢复失败，退出代码: ${RESTORE_EXIT_CODE}"
    unset PGPASSWORD
    exit ${RESTORE_EXIT_CODE}
fi

# 清理
unset PGPASSWORD

# 显示执行时间
end_timer

log_info "恢复脚本执行完成"
exit 0
